package com.zly.offer.backend.controller;

import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.zly.offer.backend.entity.AICharacter;
import com.zly.offer.backend.entity.ChatMessage;
import com.zly.offer.backend.service.CharacterService;
import com.zly.offer.backend.service.ChatMessageService;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.util.CollectionUtils;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.List;

@Slf4j
@RequestMapping("/chat")
@RestController
public class AICharacterController {
    @Autowired
    private ChatModel chatModel;
    @Resource
    private RedisTemplate<String, String> redisTemplate;
    @Autowired
    private CharacterService characterService;
    @Resource
    private ChatMessageService chatMessageService;

    @GetMapping("/character/{id}")
    public Flux<String> chat(@PathVariable Integer id,
                             @RequestParam Integer conversationId,
                             @RequestParam String question,
                             @RequestParam Integer character_id) {
//        SecurityContextHolder.getContext().getAuthentication().getName();
        String redis_key = "chat:" + id + ":" + conversationId;
        List<ChatMessage> list = chatMessageService.list(new LambdaQueryWrapper<ChatMessage>()
                .eq(ChatMessage::getUserId, id)
                .eq(ChatMessage::getConversationId, conversationId)
        );
        if (CollectionUtils.isEmpty(list)) {
            AICharacter aiCharacter = characterService.getById(character_id);
            ChatMessage chatMessage = new ChatMessage();
            chatMessage.setMessageType(MessageType.SYSTEM);
            chatMessage.setConversationId(conversationId);
            chatMessage.setUserId(id);
            chatMessage.setContent(aiCharacter.getPrompt());
            chatMessageService.save(chatMessage);//保存系统信息
            list = new ArrayList<>();
            list.add(chatMessage);
        }
        ChatMessage chatMessage = new ChatMessage();

        chatMessage.setMessageType(MessageType.USER);
        chatMessage.setConversationId(conversationId);
        chatMessage.setUserId(id);
        chatMessage.setContent(question);
        chatMessageService.save(chatMessage);//保存用户问题
        list.add(chatMessage);
        StringBuffer stringBuffer = new StringBuffer();
        ArrayList<Message> arrayList = new ArrayList<>();
        list.forEach(chat -> {
            if (chat.getMessageType() == MessageType.SYSTEM) {
                arrayList.add(new SystemMessage(chat.getContent()));
            }
            if (chat.getMessageType() == MessageType.USER) {
                arrayList.add(new SystemMessage(chat.getContent()));
            }
            if (chat.getMessageType() == MessageType.ASSISTANT) {
                arrayList.add(new AssistantMessage(chat.getContent()));
            }
        });
        return chatModel.stream(new Prompt(arrayList))
                .map(chatResponse -> chatResponse.getResult().getOutput().getText())
                .doOnEach(res -> {
                    if (res.hasValue()) {
                        stringBuffer.append(res.get());
                    }
                })
                .doOnComplete(() -> {
                    ChatMessage message = new ChatMessage();
                    message.setMessageType(MessageType.ASSISTANT);
                    message.setConversationId(conversationId);
                    message.setUserId(id);
                    message.setContent(stringBuffer.toString());
                    log.debug("message:{}", stringBuffer);
                    chatMessageService.save(message);
                });
    }
}
